import os
import argparse
import time

from piano_transcription_inference import PianoTranscription, sample_rate, load_audio


def inference(args):
    # Arugments & parameters
    audio_path = args.audio_path
    output_midi_path = audio_path.replace(".mp3",".mid")
    device =  'cpu'
 
    # Load audio
    (audio, _) = load_audio(audio_path, sr=sample_rate, mono=True)

    # Transcriptor
    transcriptor = PianoTranscription(device=device, checkpoint_path=None)
    # Transcribe and write out to MIDI file
    transcribe_time = time.time()
    transcribed_dict = transcriptor.transcribe(audio, output_midi_path)
    print('Transcribe time: {:.3f} s'.format(time.time() - transcribe_time))
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--audio_path', type=str, required=True)
    args = parser.parse_args()
    inference(args)